from ultralytics.models.yolo.detect import DetectionPredictor
from ultralytics.utils import ASSETS
import torch
import torch.distributed as dist
import argparse
from ultralytics.data import build_dataloader, build_llltData
from ultralytics import YOLO
import os
import swanlab
os.environ['WANDB_DISABLED'] = 'true'
os.environ['OMP_NUM_THREADS'] = '1'
os.environ["CUDA_VISIBLE_DEVICES"] = "5,6,7"
# os.environ['CUDA_VISIBLE_DEVICES'] = '6,7'  # 指定可见的GPU


# 初始化分布式环境
def init_distributed_mode():
    if 'RANK' in os.environ and 'WORLD_SIZE' in os.environ:
        rank = int(os.environ['RANK'])
        world_size = int(os.environ['WORLD_SIZE'])
        gpu = int(os.environ['LOCAL_RANK'])
    else:
        print('Not using distributed mode')
        return False, 0, 1, 0

    torch.cuda.set_device(gpu)
    dist_backend = 'nccl'
    dist.init_process_group(
        backend=dist_backend,
        init_method='env://',
        world_size=world_size,
        rank=rank
    )
    return True, rank, world_size, gpu


# 参数解析
parser = argparse.ArgumentParser()
parser.add_argument("--data", type=str, default="lllt.yaml", help="data yaml file")
parser.add_argument("--model", type=str, default="snn_yolov8l.yaml", help="model yaml file")
parser.add_argument("--weights", type=str, default="./runs/train/exp_llt_train9_llt_pretrain/weights/best.pt", help="weights file")
parser.add_argument("--epochs", type=int, default=100, help="epochs num")
parser.add_argument("--batch", type=int, default=4, help="batch size")
parser.add_argument("--workers", type=int, default=8, help="workers num")
args = parser.parse_args()

# 主函数


def main():
    pred_args = dict(
        model=args.weights,
        data=args.data,
        batch=args.batch,
        device=[2],
        # half=True,
        # amp=False,
        # single_cls=False,
        optimizer='auto',
        project='runs/test',
        name='exp_llt_pred2',
        exist_ok=True,
        source='/data1/lkf24/data/NESR/test/crossroad/channel2'
    )

    swanlab.init(
        project="llt_SpikeYOLO",
        workspace="Linexus",
        experiment_name=pred_args['name'],
        config=pred_args,
        logdir='./runs/log',
        load='ultralytics/cfg/default.yaml',
        mode='cloud',
        
    )

    predictor = DetectionPredictor(overrides=pred_args)
    predictor.predict_cli(isMulti=True)  # TODO 替换成你的图片路径


if __name__ == "__main__":
    main()
